import asyncio
import aiosqlite
import json
from openai import AsyncOpenAI
from pathlib import Path
from tqdm.asyncio import tqdm
import os
import re
from itertools import cycle # 导入 cycle

# --- 1. 配置 ---
# 包含 JSON 文件的文件夹 (请确保已创建此文件夹)
DATA_FOLDER = Path(r"C:\Users\Administrator\Projects\lmy\outputs\step2_sentence_extraction\t2")
# 数据库文件路径
DB_PATH = "sentiment_analysis.db"
# 并发请求数量 (这是总并发数)
CONCURRENCY_LIMIT = 24
# 积极句子的目标数量
POSITIVE_GOAL = 200000

# --- MODIFIED: 修改为模型标识符列表 ---
# LM Studio 服务器地址
LMSTUDIO_BASE_URL = "http://localhost:1234/v1"
# 在这里添加你在同一个服务上加载的所有模型标识符
MODEL_IDENTIFIERS = [
    "qwen/qwen3-4b",
    "qwen/qwen3-4b:2",
]
API_KEY = "lm-studio"  # LM Studio 的默认 key

# --- 2. Flag: 模型模式 ---
# True:  使用 temperature=0，响应更快、更确定 (推荐用于分类任务)
# False: 使用 temperature=0.7，响应更多样，可能更“思考”
USE_FAST_MODE = True

# --- 3. AI 系统提示 ---
# 保持你原来的系统提示不变
SYSTEM_PROMPT = """请你根据提供的企业气候风险披露的两大核心类别：气候风险暴露披露和气候风险防范披露，对我后续提供的句子进行分类。
标签的含义是:1.当前句子论述是否属于气候风险暴露披露，反映了企业面临气候风险？是的话则分类结果为0；
2.当前句子论述是否属于气候风险防范披露，反映了企业采取措施对气候风险进行防范？是的话则分类结果则为1，
3.若既不是气候风险暴露披露，也不是气候风险防范披露则分类结果为2。
注意你的输出格式：用一两句话简述你的思考过程。例如，分析句子的关键词、语气和整体含义，并用<thought>和</thought>标签对包裹你的思考过程简述。
另外，用<result>和</result>标签对包裹你的分类结果，然后终止思考和输出"""

async def setup_database():
    """
    初始化SQLite数据库，并创建包含详细思考和摘要思考的列。
    """
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS sentiments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                sentence TEXT NOT NULL UNIQUE,
                result INTEGER NOT NULL,
                detailed_thought TEXT,
                summary_thought TEXT
            )
        """)
        await db.commit()
        # 注意：由于任务目标已变为分类，这里的“目标计数”逻辑可能需要调整。
        # 我们暂时仍然计算结果不为2（非其他类）的数量作为进度。
        cursor = await db.execute("SELECT COUNT(id) FROM sentiments WHERE result != 2")
        count = (await cursor.fetchone())[0]
        return count


async def get_all_sentences(folder_path: Path):
    """递归查找所有 JSON 文件并提取所有句子。"""
    print(f"🔍 正在 '{folder_path}' 文件夹中搜索 JSON 文件...")
    json_files = list(folder_path.rglob("*.json"))
    if not json_files:
        print(f"⚠️ 在 {folder_path} 中未找到任何 JSON 文件。请创建该文件夹并添加文件。")
        return

    # 获取数据库中所有已存在的句子，避免重复处理
    existing_sentences = set()
    async with aiosqlite.connect(DB_PATH) as db:
        async with db.execute("SELECT sentence FROM sentiments") as cursor:
            async for row in cursor:
                existing_sentences.add(row[0])
    print(f"数据库中已存在 {len(existing_sentences)} 个句子，本次运行将跳过它们。")

    for file_path in json_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                if isinstance(data, list):
                    for sentence in data:
                        if isinstance(sentence, str) and sentence not in existing_sentences:
                            yield sentence
        except (json.JSONDecodeError, IOError) as e:
            print(f"❌ 无法读取或解析文件 {file_path}: {e}")


# --- MODIFIED: 添加 model_identifier 参数 ---
async def analyze_sentence(client: AsyncOpenAI, sentence: str, semaphore: asyncio.Semaphore, model_identifier: str):
    """
    向LLM发送句子进行分析，并解析返回的三个标签内容。
    返回一个元组 (sentence, result, detailed_thought, summary_thought)，如果失败则返回 None。
    """
    async with semaphore:
        # 新的输出格式内容更多，需要增加max_tokens
        max_tokens = 1024
        temperature = 0.0 if USE_FAST_MODE else 0.2

        try:
            response = await client.chat.completions.create(
                model=model_identifier, # <-- MODIFIED: 使用传入的模型标识符
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": f'句子: "{sentence}"'}
                ],
                max_tokens=max_tokens,
                temperature=temperature,
            )
            content = response.choices[0].message.content.strip()

            # 保持你原来的解析逻辑不变
            think_match = re.search(r'<think>(.*?)</think>', content, re.DOTALL)
            detailed_thought = think_match.group(1).strip() if think_match else "模型未提供<think>标签。"
            thought_match = re.search(r'<thought>(.*?)</thought>', content, re.DOTALL)

            summary_thought = thought_match.group(1).strip() if thought_match else "模型未提供<thought>标签。"
            result_match = re.search(r'<result>(.*?)</result>', content, re.DOTALL)
            if result_match:
                try:
                    result = int(result_match.group(1).strip())
                    return (sentence, result, detailed_thought, summary_thought)
                except (ValueError, IndexError):
                    tqdm.write(f"⚠️ 模型在<result>标签内返回了非数字内容: '{result_match.group(1)}'")
            else:
                tqdm.write(f"⚠️ 模型未提供<result>标签。完整响应: '{content[:200]}...'")
            return None


        except Exception as e:
            tqdm.write(f"❌ 处理句子时发生错误 (模型: {model_identifier}) '{sentence}': {e}")
            return None


async def main():
    """主函数，负责协调整个并发处理流程。"""
    print("--- LM Studio 并发情感分析系统 ---")
    print(f"🚀 并发数: {CONCURRENCY_LIMIT}")
    print(f"⚡ 快速模式: {'启用' if USE_FAST_MODE else '关闭'}")
    print(f"🎯 目标: {POSITIVE_GOAL} 个积极句子")
    print(f"🤖 将轮流调用模型: {MODEL_IDENTIFIERS}")
    print("-" * 40)

    if not DATA_FOLDER.exists():
        DATA_FOLDER.mkdir()
        print(f"已创建数据文件夹: '{DATA_FOLDER}'。请将您的 JSON 文件放入其中。")
        return

    positive_count = await setup_database()
    if positive_count >= POSITIVE_GOAL:
        print(f"✅ 目标已达成！数据库中已有 {positive_count} 个积极句子。")
        return

    client = AsyncOpenAI(base_url=LMSTUDIO_BASE_URL, api_key=API_KEY)
    semaphore = asyncio.Semaphore(CONCURRENCY_LIMIT)
    tasks = []

    # --- MODIFIED: 创建模型标识符的循环器 ---
    model_cycler = cycle(MODEL_IDENTIFIERS)

    # 使用 aiosqlite.connect 作为上下文管理器，确保连接被正确关闭
    async with aiosqlite.connect(DB_PATH) as db:
        # 初始化进度条
        pbar = tqdm(total=POSITIVE_GOAL, initial=positive_count, desc="积极句子进度", unit="句")
        processed_counter = 0

        # 从生成器创建所有分析任务
        async for sentence in get_all_sentences(DATA_FOLDER):
            if len(sentence) > 400:
                continue
            if positive_count < POSITIVE_GOAL:
                # --- MODIFIED: 轮询选择模型并创建任务 ---
                model_id_to_use = next(model_cycler)
                task = asyncio.create_task(analyze_sentence(client, sentence, semaphore, model_id_to_use))
                tasks.append(task)
            else:
                break  # 如果一开始就达到目标，则不创建任何任务

        # 当任务完成时立即处理
        for future in asyncio.as_completed(tasks):
            if positive_count >= POSITIVE_GOAL:
                break  # 达到目标后，停止等待新结果

            result = await future
            processed_counter += 1
            pbar.set_postfix({"已处理句子总数": processed_counter})

            if result:
                # 保持你原来的数据库插入逻辑
                sentence, result_val, detailed_thought, summary_thought = result
                try:
                    await db.execute(
                        "INSERT OR IGNORE INTO sentiments (sentence, result, detailed_thought, summary_thought) VALUES (?, ?, ?, ?)",
                        (sentence, result_val, detailed_thought, summary_thought)
                    )
                    await db.commit()

                    if result_val != 2:
                        positive_count += 1
                        pbar.update(1)

                except aiosqlite.IntegrityError:
                    pbar.write(f"Info: Sentence '{sentence[:50]}...' was already in the database.")

        # 任务循环结束后，取消所有仍在运行的剩余任务
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)  # 等待取消完成

        pbar.close()

    print("\n" + "-" * 40)
    if positive_count >= POSITIVE_GOAL:
        print(f"✅ 成功! 已达成 {POSITIVE_GOAL} 个积极句子的目标。")
    else:
        print(f"⏹️ 流程结束。共找到 {positive_count} 个积极句子，未达到 {POSITIVE_GOAL} 的目标。")
    print(f"本次运行共处理了 {processed_counter} 个句子。")
    print(f"数据库已保存至: {DB_PATH}")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n检测到用户中断，正在退出...")

